package com.jt1078.streamserver.flv;

import java.io.ByteArrayOutputStream;
import java.io.IOException;

public class FLVMuxer {
    public static byte[] generateFLVHeader() {
        return new byte[] {'F', 'L', 'V', 0x01, 0x01, 0,0,0,9, 0,0,0,0};
    }

    public static byte[] removeStartCode(byte[] nalu) {
        int offset = 0;
        if (nalu.length > 4 && nalu[0] == 0 && nalu[1] == 0 && nalu[2] == 0 && nalu[3] == 1)
            offset = 4;
        else if (nalu.length > 3 && nalu[0] == 0 && nalu[1] == 0 && nalu[2] == 1)
            offset = 3;
        byte[] raw = new byte[nalu.length - offset];
        System.arraycopy(nalu, offset, raw, 0, raw.length);
        return raw;
    }

    public static byte[] createAvcSeqHeader(byte[] sps, byte[] pps) throws IOException {
        // SPS和PPS都要去掉起始码
        sps = removeStartCode(sps);
        pps = removeStartCode(pps);

        ByteArrayOutputStream out = new ByteArrayOutputStream();
        out.write(0x09); // TagType: Video
        int dataSize = 16 + sps.length + 3 + pps.length;
        out.write((dataSize >> 16) & 0xFF);
        out.write((dataSize >> 8) & 0xFF);
        out.write(dataSize & 0xFF);
        out.write(0); out.write(0); out.write(0); out.write(0); // ts/tsExt
        out.write(0); out.write(0); out.write(0); // StreamID
        out.write(0x17); // KeyFrame + AVC
        out.write(0x00); // AVC sequence header
        out.write(0); out.write(0); out.write(0); // CompositionTime
        out.write(0x01); // version
        out.write(sps[1]); // profile
        out.write(sps[2]); // compat
        out.write(sps[3]); // level
        out.write(0xFF); // nal size=4
        out.write(0xE1); // num of sps
        out.write((sps.length >> 8) & 0xFF); out.write(sps.length & 0xFF); out.write(sps);
        out.write(1); // num of pps
        out.write((pps.length >> 8) & 0xFF); out.write(pps.length & 0xFF); out.write(pps);
        int total = out.size();
        out.write((total >> 24) & 0xFF); out.write((total >> 16) & 0xFF); out.write((total >> 8) & 0xFF); out.write(total & 0xFF);
        return out.toByteArray();
    }

    public static byte[] createVideoTag(byte[] nalu, boolean isKeyframe, int ts) throws IOException {
        nalu = removeStartCode(nalu);
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        out.write(0x09);
        int bodySize = 5 + 4 + nalu.length;
        out.write((bodySize >> 16) & 0xFF);
        out.write((bodySize >> 8) & 0xFF);
        out.write(bodySize & 0xFF);
        out.write((ts >> 16) & 0xFF); out.write((ts >> 8) & 0xFF); out.write(ts & 0xFF); out.write(0);
        out.write(0); out.write(0); out.write(0);
        out.write(isKeyframe ? 0x17 : 0x27); // Keyframe/Interframe + AVC
        out.write(0x01); // AVC NALU
        out.write(0); out.write(0); out.write(0);
        // 4字节nalu长度
        out.write((nalu.length >> 24) & 0xFF);
        out.write((nalu.length >> 16) & 0xFF);
        out.write((nalu.length >> 8) & 0xFF);
        out.write(nalu.length & 0xFF);
        out.write(nalu);
        int total = out.size();
        out.write((total >> 24) & 0xFF); out.write((total >> 16) & 0xFF); out.write((total >> 8) & 0xFF); out.write(total & 0xFF);
        return out.toByteArray();
    }
}
